#pragma once

#include <vector>
#include <string>

using std::string;
using std::vector;

constexpr uint64_t SSDSIZE  = 200000; // 200 GiB SSD
constexpr uint64_t NVMSIZE  = 5120;   // 5 GiB NVM
constexpr uint64_t DRAMSIZE = 1920;   // 2 GiB DRAM (plus memtables)
constexpr   uint64_t DEFAULT_VAL_SIZE = 1008; // 1KiB key value item
constexpr uint64_t MEMTABLE = 64;     // 64 MiB memtable

enum {OP_INSERT, OP_READ, OP_UPDATE, OP_REMOVE, OP_SCAN};

struct ParamType {
    bool write_sync;
    int val_size;
};

using _key_t = uint64_t;

class DB {
 public:
    virtual void insert(_key_t key, std::string value) = 0;

    virtual bool find(_key_t key, std::string *value) = 0;

    virtual bool update(_key_t key, std::string value) = 0;

    virtual bool remove(_key_t key) = 0;

    virtual int scan(_key_t key, int range) = 0;

    virtual ~DB() {};
};

/////////////////////////////////////////////////////////////////////
// DNSKV: KV Cache in DRAM, B+-tree in NVM, Slab file & Log file in SSD
/////////////////////////////////////////////////////////////////////
#include "dnskv/dnskv.h"
class DNSKVDB : public DB
{
public:
    DNSKVDB(ParamType p) {
        dnskv::DNSOpt opt = {
            .dbdir  = "/data/nvme0/dnskv",
            .nvmdir = "/mnt/pmem/dnskv",
            .ssdsize  = SSDSIZE * 1024 * 1024,
            .nvmsize  = NVMSIZE * 1024 * 1024,
            .dramsize = (DRAMSIZE + MEMTABLE * 2) * 1024 * 1024,
            .avgsize  = p.val_size,
        };
        db_ = new dnskv::DNSKV(opt);
        sync_ = p.write_sync;
    }

    ~DNSKVDB() {
        delete db_;
    }

    void insert(_key_t key, std::string value) {
        db_->Put(key, value, sync_);
    }

    bool find(_key_t key, std::string *value) {
        return db_->Get(key, *value);
    }

    bool update(_key_t key, std::string value) {
        return db_->Update(key, value, sync_);
    }

    int scan(_key_t key, int range) {
        std::vector<string> res;
        auto ret = db_->Scan(key, range, res);
        return ret;
    }

    bool remove(_key_t key) {
        return db_->Delete(key, sync_);
    }

private:
    dnskv::DNSKV * db_;
    bool sync_;
};


/////////////////////////////////////////////////////////////////////
// MatrixKV: LSM-tree based KVStore and L0 in NVM
/////////////////////////////////////////////////////////////////////
#include "rocksdb/db.h"
#include "rocksdb/slice.h"
#include "rocksdb/options.h"
#include "rocksdb/nvm_option.h"
#include "rocksdb/table.h"
class MatrixKVDB : public DB
{
public:
    MatrixKVDB(ParamType p) {
        rocksdb::Options options;
        options.OptimizeLevelStyleCompaction();
        std::shared_ptr<rocksdb::Cache> cache = rocksdb::NewLRUCache(DRAMSIZE * 1024 * 1024);
        rocksdb::BlockBasedTableOptions table_options;
        table_options.block_cache = cache;
        options.table_factory.reset(NewBlockBasedTableFactory(table_options));

        options.create_if_missing = true;
        options.nvm_setup = std::make_shared<rocksdb::NvmSetup>();
        options.nvm_setup->Level0_column_compaction_trigger_size  = (NVMSIZE * 7 / 8) * 1024 * 1024;
        options.nvm_setup->Level0_column_compaction_slowdown_size = (NVMSIZE * 15 / 16) * 1024 * 1024;
        options.nvm_setup->Level0_column_compaction_stop_size     = NVMSIZE * 1024 * 1024;
        options.nvm_setup->pmem_path = "/mnt/pmem/matrixkv";
        options.nvm_setup->use_nvm_module = true;
        options.compression = rocksdb::kNoCompression;
        std::string dbpath = "/data/nvme0/matrixkv";
        options.max_bytes_for_level_base = options.nvm_setup->Level0_column_compaction_stop_size;
        
        // 打开数据库 
        rocksdb::Status s = rocksdb::DB::Open(options, dbpath, &db_);
        assert(s.ok());

        wop_ = rocksdb::WriteOptions();
        wop_.sync = p.write_sync;
        rop_ = rocksdb::ReadOptions();
    }

    ~MatrixKVDB () {
        delete db_;
    }

    void insert(_key_t key, std::string value) {
        rocksdb::Status s = db_->Put(wop_, std::to_string(key), value);
    }

    bool find(_key_t key, std::string *value) {
        rocksdb::Status s = db_->Get(rop_, std::to_string(key), value);
        return !s.IsNotFound();
    }

    bool update(_key_t key, std::string value) {
        rocksdb::Status s = db_->Put(wop_, std::to_string(key), value);
        return !s.IsNotFound();
    }

    int scan(_key_t key, int range) {
        auto iter = db_->NewIterator(rop_);
        iter->Seek(rocksdb::Slice(std::to_string(key)));
        
        vector<string> res;
        for(int i = 0; i < range; i++) {
            res.push_back(iter->value().ToString());
            iter->Next();
        }

        return res.size();
    }

    bool remove(_key_t key) {
        rocksdb::Status s = db_->Delete(wop_, std::to_string(key));
        return !s.IsNotFound();
    }

private:
    rocksdb::DB* db_;
    rocksdb::WriteOptions wop_;
    rocksdb::ReadOptions rop_;
};

/////////////////////////////////////////////////////////////////////
// MySLM: Btree in NVM and KV log in SSD
/////////////////////////////////////////////////////////////////////
#include "myslm/myslm.h"
class MySLMDB : public DB
{
public:
    MySLMDB(ParamType p) {
        myslm::SLMOpt opt = {
            .dbdir  = "/data/nvme0/myslm",
            .nvmdir = "/mnt/pmem/myslm",
            .nvmsize  = NVMSIZE * 1024 * 1024,
        };
        db_ = new myslm::MySLM(opt);
        sync_ = p.write_sync;
    }

    ~MySLMDB() {
        delete db_;
    }

    void insert(_key_t key, std::string value) {
        db_->Put(key, value, sync_);
    }

    bool find(_key_t key, std::string *value) {
        return db_->Get(key, *value);
    }

    bool update(_key_t key, std::string value) {
        return db_->Update(key, value, sync_);
    }

    int scan(_key_t key, int range) {
        std::vector<string> res;
        auto ret = db_->Scan(key, range, res);
        return ret;
    }

    bool remove(_key_t key) {
        return db_->Delete(key, sync_);
    }

private:
    myslm::MySLM * db_;
    bool sync_;
};


/////////////////////////////////////////////////////////////////////
// B+-tree: B+-tree inner nodes in NVM and B+-tree leaf nodes in SSD
/////////////////////////////////////////////////////////////////////
#include "btree/nsbtree.h"
class BtreeDB : public DB
{
public:
    BtreeDB(ParamType p) {
        nsbtree::NSBOpt opt = {
            .dbdir  = "/data/nvme0/btree",
            .nvmdir = "/mnt/pmem/btree",
            .dramsize = (DRAMSIZE + MEMTABLE * 2) * 1024 * 1024,
            .nvmsize  = NVMSIZE * 1024 * 1024,
            .ssdsize  = SSDSIZE * 1024 * 1024 * 2,
        };
        db_ = new nsbtree::NSBtree(opt);
        sync_ = p.write_sync;
    }

    ~BtreeDB() {
        delete db_;
    }

    void insert(_key_t key, std::string value) {
        db_->Put(key, value, sync_);
    }

    bool find(_key_t key, std::string *value) {
        return db_->Get(key, *value);
    }

    bool update(_key_t key, std::string value) {
        return db_->Update(key, value, sync_);
    }

    int scan(_key_t key, int range) {
        std::vector<string> res;
        auto ret = db_->Scan(key, range, res);
        return ret;
    }

    bool remove(_key_t key) {
        return db_->Delete(key, sync_);
    }

private:
    nsbtree::NSBtree * db_;
    bool sync_;
};

//==============================================================
// GET INSTANCE
//==============================================================
inline DB *getInstance(const std::string type, ParamType p) {
  if (type == "dnskv")
    return new DNSKVDB(p);
  else if(type == "matrixkv")
    return new MatrixKVDB(p);
  else if(type == "myslm")
    return new MySLMDB(p);
  else if(type == "btree")
    return new BtreeDB(p);
  else {
    fprintf(stderr, "Unknown index type : %s\n", type.c_str());
    exit(1);
  }
  
  return nullptr;
}